import json

from sparkai.llm.llm import ChatSparkLLM, ChunkPrintHandler
from sparkai.core.messages import ChatMessage

#星火认知大模型v3.5的URL值，其他版本大模型URL值请前往文档（https://www.xfyun.cn/doc/spark/Web.html）查看
SPARKAI_URL = 'wss://spark-api.xf-yun.com/v3.5/chat'
#星火认知大模型调用秘钥信息，请前往讯飞开放平台控制台（https://console.xfyun.cn/services/bm35）查看
SPARKAI_APP_ID = 'd361809e'
SPARKAI_API_SECRET = 'Y2MxNTdhYzI1M2IyNGVjYjIzYjM3ZThk'
SPARKAI_API_KEY = 'e03f74be5d4693b3821f43bcc9193720'
#星火认知大模型v3.5的domain值，其他版本大模型domain值请前往文档（https://www.xfyun.cn/doc/spark/Web.html）查看
SPARKAI_DOMAIN = 'generalv3.5'


def readtxt():
    # 假设文件名为 example.txt，位于当前目录下
    file_path = 'app/brochure.txt'

    # 使用 with 语句打开文件，这样可以确保文件在使用后会被正确关闭
    with open(file_path, 'r', encoding='utf-8') as file:
        # 读取文件全部内容
        content = file.read()
        # 打印内容
    return content


def starfire(message):
    init = readtxt()
    end = '请你根据我所提供的关于联邦学习的内容，回答我的问题'
    spark = ChatSparkLLM(
        spark_api_url=SPARKAI_URL,
        spark_app_id=SPARKAI_APP_ID,
        spark_api_key=SPARKAI_API_KEY,
        spark_api_secret=SPARKAI_API_SECRET,
        spark_llm_domain=SPARKAI_DOMAIN,
        streaming=False,
    )
    messages = [ChatMessage(
        role="user",
        content=end + init + message
    )]
    handler = ChunkPrintHandler()
    a = spark.generate([messages], callbacks=[handler])
    return a.generations


if __name__ == '__main__':
    spark = ChatSparkLLM(
        spark_api_url=SPARKAI_URL,
        spark_app_id=SPARKAI_APP_ID,
        spark_api_key=SPARKAI_API_KEY,
        spark_api_secret=SPARKAI_API_SECRET,
        spark_llm_domain=SPARKAI_DOMAIN,
        streaming=False,
    )
    messages = [ChatMessage(
        role="user",
        content='你好呀'
    )]
    handler = ChunkPrintHandler()
    a = spark.generate([messages], callbacks=[handler])
    for chat_generation in a.generations:
        content = chat_generation[0].message.content
        print(content)
